#define _GNU_SOURCE 1
#include <sched.h> //for binding process to processor

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>
#include <sys/un.h>
#include <netinet/in.h>
#include <sys/wait.h>
#include <sys/time.h>
#include <time.h>
#include <pthread.h>



#define BODY_SIZE 2
#define ITERATIONS 100

char message_id[INET_ADDRSTRLEN];
char message_body[BODY_SIZE];



char** srcIpsSpoof =  NULL;
unsigned long numIps = 0;
int client_time = 10;
char sock_path[108];

/**
 * reads the next ip from a given file
 * @param file the file to read from
 * @param ip where to store the read ip
 * @return 1 on success, 0 otherwise
 */
int getIp(FILE *file, char* ip) {
	if(fscanf(file, "%s", ip) == 1){
		return 1;
	}
	return 0;
}
/**
 * reads ips from file and stores them into a local data structure
 */
void readIps(char* filename) {
	FILE *file;
	if ( (file = fopen(filename, "r")) == NULL) {
		fprintf(stderr,"File %s could not be opened.\n", filename);
		exit(EXIT_FAILURE);
	}
	fprintf(stderr,"reading file...");
	fflush(stdout);
	numIps = 0;
	char srcIp[INET_ADDRSTRLEN];
	while ((getIp(file, &srcIp)) && !feof(file)) {
		numIps++;
	}
	rewind(file);
	srcIpsSpoof = (char**)malloc(numIps * sizeof(char*));
	unsigned long cnt = 0;
	while (!feof(file) && (getIp(file, &srcIp))) {
		char* entry = (char*)malloc(INET_ADDRSTRLEN);
		strcpy(entry, srcIp);
		srcIpsSpoof[cnt++] = entry;
		//printf("ip: %s\n", entry);
	}
	fclose(file);
	fprintf(stderr,"done, number of ips : %ld \n", numIps);
}



void stress_balance(){
	struct sockaddr_un balance_address;
	balance_address.sun_family = AF_UNIX;
	strcpy(balance_address.sun_path, sock_path);
	int len = sizeof(balance_address);


	memset(message_body, 'a', BODY_SIZE);
	int message_body_size = sizeof(message_body);
	int message_id_size = sizeof(message_id);

	srand(time(NULL)*getpid());
	struct timeval startTime, endTime;
	gettimeofday(&startTime, NULL);
	int stop = 0;
	long request_num = 0;
	unsigned long nfailures = 0;
	unsigned long tsent = 0;
	unsigned long tresv = 0;
	while(!stop){
		int client_socket;
		int nsent = 0;

		if ((client_socket = socket(AF_UNIX, SOCK_STREAM, 0)) == -1) {
			perror("socket error");
			exit(1);
		}

		if (connect(client_socket, (struct sockaddr *)&balance_address, len) == -1) {
			perror("connect error");
			exit(1);
		}


		char* src_ip = srcIpsSpoof[rand()%numIps];
		strcpy(message_id, src_ip);

		if ((nsent = send(client_socket, &message_id, message_id_size, 0)) == -1) {
			perror("send error");
			exit(1);
		}
		tsent+=nsent;

		int p= 0;
		for(p = 0; p < ITERATIONS; p++){
			if ((nsent = send(client_socket, &message_body, message_body_size, 0)) == -1) {
				perror("send error");
				exit(1);
			}
			tsent+=nsent;
			char server_response[BODY_SIZE + 1];
			int nrecv = -1;
			if ((nrecv=recv(client_socket, server_response, BODY_SIZE, 0)) > 0) {
				server_response[nrecv] = '\0';
				//printf("> %s\n", server_response);
				tresv+=nrecv;
			} else {
				nfailures++;
				if (nrecv < 0){
				//	perror("recv");
				}
				else{
					//printf("Server closed connection, current client %d \n", i);
				}
				break;
			}
		        gettimeofday(&endTime, NULL);
                	int delay = (endTime.tv_sec - startTime.tv_sec);
	                if(delay > client_time){
				break;
        	                stop = 1;
                	}
		}

		gettimeofday(&endTime, NULL);
		int delay = (endTime.tv_sec - startTime.tv_sec);
		if(delay > client_time){
			stop = 1;
		}
		request_num++;
		close(client_socket);
	}
	fprintf(stderr,"\n%d:requests:tsent:tresv:failures %ld:%ld:%ld:%ld\n", getpid(), request_num, tsent, tresv, nfailures);

	pthread_exit(0);
}

int main(int argc, char *argv[])
{
	int pid = getpid();
	if(argc < 7){
		fprintf(stderr,"Usage: balance-client <number of clients-process> <number of client threads> <balance socket path> <ips file path> <time to work> [numOfCPUS:commaSeparetedCPUs]");
	}

	int number_clients = atoi(argv[1]);
	fprintf(stderr,"# clients %d \n",number_clients);

	int number_clients_threads = atoi(argv[2]);
	fprintf(stderr,"# clients threads %d \n",number_clients_threads);


	strcpy(sock_path, argv[3]);

	readIps(argv[4]);

	client_time = atoi(argv[5]);
	fprintf(stderr,"clients will run for %d seconds \n",client_time);


	/*bind to specific processor*/
	char*  cpu_arg = argv[6];
	if (cpu_arg != NULL) {
		fprintf(stderr,"cpu arg %s\n", cpu_arg);
		cpu_set_t* cpu_set;
		char* p = NULL;
		int num_of_cpus = 0;
		if ((p = strtok(cpu_arg, ":")) == NULL) {
			perror("split error");
			return;
		} else {
			num_of_cpus = atoi(p);
			fprintf(stderr,"num of cpus %d\n", num_of_cpus);
			cpu_set = CPU_ALLOC(num_of_cpus);
		}
		size_t size = CPU_ALLOC_SIZE(num_of_cpus);

		CPU_ZERO_S(size, cpu_set);
		while ((p = strtok(NULL, ",")) != NULL) {
			int cpu_num = atoi(p);
			fprintf(stderr,"cpu_num to run on %d\n", cpu_num);
			CPU_SET_S(cpu_num, size, cpu_set);
		}

		fflush(stderr);

		if (sched_setaffinity(pid, size, cpu_set) < 0) {
			perror("sched_setaffinity");
		}

		if (sched_getaffinity(pid, size, cpu_set) < 0) {
			perror("sched_getaffinity");
			return -1;
		}
		fprintf(stderr, "pid %d's affinity mask: %d\n", pid, *cpu_set);
		CPU_FREE(cpu_set);
	}

	signal(SIGCHLD, SIG_IGN);

	int process_index = 0;
	for(process_index = 0; process_index < number_clients; process_index++){
		if (fork() == 0) {//child
			fprintf(stderr,"%d : client process %d created\n",pid, process_index);
			pthread_t* threads = malloc(number_clients_threads*sizeof(pthread_t));
			pthread_attr_t attr;
			pthread_attr_init(&attr);
			pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);
			int i = 0;
			for(i = 0; i < number_clients_threads; i++){
				if (pthread_create(&(threads[i]), &attr, stress_balance, NULL) != 0) {
					perror("Error creating thread\n");
					exit(-1);
				}
			}

			void * status;
			for (i = 0; i < number_clients_threads; i++) {
				if (pthread_join(threads[i], &status) != 0) {
					perror("Error joining thread\n");
					exit(-1);
				}
			}
			pthread_attr_destroy(&attr);
			free(threads);
			return 0;
		} else {//parent
			continue;
		}
	}

	return 0;
}
